package com.tuniu.agents.impl.base;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.agent.*;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.command.CommandManagerUtil;
import com.tuniu.agents.common.enums.MarkdownActionEnum;
import com.tuniu.agents.common.message.BaseAgentNewResp;
import com.tuniu.agents.common.message.RecQuestionTypeEnum;
import com.tuniu.agents.common.message.SuggestionQuestion;
import com.tuniu.agents.common.message.adaper.AdapterRecQuestion;
import com.tuniu.agents.common.util.MarkdownActionUtil;
import com.tuniu.agents.common.util.NumberingUtil;
import com.tuniu.agents.message.MessageRender;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.resolution.ToolCallbackResolver;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;

import java.util.*;
import java.util.stream.Collectors;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * Question Agent
 */
@Slf4j
public class Analyst extends OptionsAgent {

    public static final String AGENT_ID = "Anylyst";

    public static final String PROMPT_ID = "analyst";

    private final AgentManager agentManager;

    @Autowired
    private ToolCallbackResolver toolCallbackResolver;

    private final List<String> matchToolNames;

    private final Map<String, String> SelectionAgents;

    private MessageRender messageRender;

    /**
     * {{user_message_current}} is assembled in Selection before coming here
     * Analyst:
     * User's previous requirements:
     * 1: XXX
     * 2: XXX
     * Current user requirement:
     * XXX
     * Current recommendation result:
     * XXX (MD format)
     */
    public static final String MESSAGE_TEMPLATE = """
            {{user_message_his}}
            Current user requirement:
            {{user_message_current}}
            """;

    public Analyst(AgentManagerGroup agentManagerGroup, List<String> matchToolNames, Map<String, String> SelectionAgents, ChatOptions chatOptions) {
        this(agentManagerGroup, matchToolNames, SelectionAgents,  AGENT_ID, PROMPT_ID,chatOptions);
    }

    public Analyst(AgentManagerGroup agentManagerGroup, List<String> matchToolNames, Map<String, String> SelectionAgents,  String agentId, String promptTemplateId, ChatOptions chatOptions) {
        super(agentManagerGroup.getModelBuilder(),
                agentManagerGroup.getChatMemory(),
                agentManagerGroup.getAdvisors().stream().map(advisor -> advisor instanceof AgentChatMemoryAdvisor ? new AgentChatMemoryAdvisor(agentManagerGroup.getChatMemory(), false, false) : advisor).collect(Collectors.toList()),
                agentManagerGroup.getPromptManager(),
                agentManagerGroup.getChatEnvironmentMemory(),
                new AgentOptions(agentId, promptTemplateId, chatOptions));
        this.agentManager = agentManagerGroup.getAgentManager();
        this.matchToolNames = matchToolNames;
        this.SelectionAgents = SelectionAgents;
        this.messageRender = agentManagerGroup.getMessageRender();
    }

    public List<AdapterRecQuestion> callQuestion(String conversationId, String message) {
        /**
         * 1. Get potential requirements and uncalled tool list from global context
         * 2. Get user information from global context, such as departure city, destination city, departure date
         * 3. Combine prompt words into the prompt
         */
        Collection<String> potential_intents_history = (Collection<String>) super.chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.POTENTIAL_INTENTS_HISTORY); // Historical potential intents
        final Collection<String> main_intent_history = (Collection<String>) Optional.ofNullable(super.chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.MAIN_INTENT_HISTORY)).orElse(List.of());; // Historical main intents
        String lastAgentId = Optional.ofNullable(super.chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.LAST_SELECTION_AGENT_ID)).map(Object::toString).orElse(null); // Last used AgentId
        String lastToolId = Optional.ofNullable(super.chatEnvironmentMemory.get(conversationId, AgentsConstants.AgentsConversationEnvironmentKey.LAST_SELECTION_TOOL_ID)).map(Object::toString).orElse(null); // Last used AgentId
        // No cross-recommendation when querying details
        Map<String, String> toolMapping = new HashMap<>();
        /*
         Agent association, key: agentId/toolName
         value: agentId
         */
        Map<String, String> agentMapping = new HashMap<>();
        int index = 1;
        StringBuffer agentSb = new StringBuffer();
        for(Map.Entry<String, String> entry : SelectionAgents.entrySet()) {
            agentMapping.put(entry.getKey(), entry.getKey());
            if(main_intent_history != null && !main_intent_history.contains(entry.getKey())) {
                agentSb.append(index++).append(".").append(entry.getKey()).append("\n");
            }
            // Set up tools
            Agent agent = agentManager.getAgent(entry.getKey());
            if (agent instanceof OptionsAgent) {
                OptionsAgent lastOptionAgent = (OptionsAgent) agent;
                if(CollectionUtils.isEmpty(lastOptionAgent.getAgentOptions().getTools())){
                    continue;
                }
                Collection<String> tools = new ArrayList<>(lastOptionAgent.getAgentOptions().getTools());
                if (CollectionUtils.isNotEmpty(tools) && lastToolId != null) { // Remove the last used tool
                    tools.remove(lastToolId);
                }
                tools = CollectionUtils.intersection(tools, matchToolNames);
                for (String toolName : tools) {
                    try {
                        FunctionCallback functionCallback = toolCallbackResolver.resolve(toolName);
                        toolMapping.put(toolName, functionCallback.getDescription());
                        agentMapping.put(toolName, entry.getKey());
                    } catch (Exception e) {
                        log.error("Tool {}", toolName, e);
                    }
                }
            }
        }
        index = 1;
        StringBuffer toolSb = new StringBuffer();
        for(Map.Entry<String, String> entry : toolMapping.entrySet()) {
            toolSb.append(index++).append(".").append(entry.getKey()).append("\n")
                    .append("\"description\":\"").append(entry.getValue()).append("\"\n");
        }
        log.info(agentSb.toString());
        log.info(toolSb.toString());
        if (StringUtils.isNotBlank(message)) {
            List<Analyst.Question> questions = chatClient.prompt()
                    .system(s -> s.params(chatEnvironmentMemory.get(conversationId)))
                    .messages(List.of(new UserMessage(message)))
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId).
                            param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, agentOptions.getAdvisorParam().get(CHAT_MEMORY_RETRIEVE_SIZE_KEY)).
                            param(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY, this.agentId).
                            param(AgentsConstants.AgentsAdvisorKey.USER_MESSAGE_TEMPLATE, MESSAGE_TEMPLATE).
                            param("agents_list", agentSb.toString()).
                            param("tool_list", toolSb.toString()))
                    .call()
                    .entity(new ParameterizedTypeReference<List<Analyst.Question>>() {
                    });
            List<SuggestionQuestion> result = questions.stream()  // Only take the first three
                    .map(question -> new SuggestionQuestion(question.result, agentMapping.getOrDefault(question.type, ""), toolMapping.containsKey(question.type) ? question.type: "")).collect(Collectors.toList());
            List<AdapterRecQuestion> recResult = questions.stream()
                    .map(question -> {
                        // Use freemarker and MarkdownUtil to beautify the return result
                        Long number = NumberingUtil.incrWithExpire(conversationId);
                        String recQuestionResult = messageRender.renderTemplate("common/analyst.ftlh", Map.of("recQuestion", question.result, NumberingUtil.KEY_NUMBER, number));
                        String recQuestionMarkdown = MarkdownActionUtil.buildMarkdownAction(recQuestionResult, MarkdownActionEnum.EVENT_SEND, "");
                        AdapterRecQuestion recQuestion = BaseAgentNewResp.sendRecQuestion("", RecQuestionTypeEnum.REC_QUESTION_SEND, recQuestionMarkdown);
                        CommandManagerUtil.saveStringAgentCommand(conversationId, "AGENT_ID" + number, number, Dispatcher.AGENT_ID, recQuestionMarkdown);
                        return recQuestion;
                    }).collect(Collectors.toList());
            return recResult.stream()
                    .limit(3) // Take the first three records
                    .collect(Collectors.toList());
        } else { // If input is empty, no recommendation
            return null;
        }
    }

    private void extracted(String lastAgentId, String lastToolId, Set<String> input, Map<String, String> toolMapping, Map<String, String> agentMapping) {
        Agent lastAgent = agentManager.getAgent(lastAgentId);
        if (lastAgent instanceof OptionsAgent) {
            OptionsAgent lastOptionAgent = (OptionsAgent) lastAgent;
            Collection<String> tools = new ArrayList<>(lastOptionAgent.getAgentOptions().getTools());
            if (CollectionUtils.isNotEmpty(tools) && lastToolId != null) { // Remove the last used tool
                tools.remove(lastToolId);
            }
            tools = CollectionUtils.intersection(tools, matchToolNames);
            for (String toolName : tools) {
                try {
                    FunctionCallback functionCallback = toolCallbackResolver.resolve(toolName);
                    input.add(functionCallback.getDescription());
                    toolMapping.put(toolName, functionCallback.getDescription());
                    agentMapping.put(functionCallback.getDescription(), lastAgentId);
                } catch (Exception e) {
                    log.error("Tool {}", toolName, e);
                }
            }
        }
    }


    /**
     * Dispatcher Bean
     *
     * @param type   Question type
     * @param result Question result
     */
    record Question(String type, String result) {


    }


}
